import collections
import re

import numpy as np
import torch
from d2l import torch as d2l


# 马尔科夫假设：一个词的出现只与前面N个词相关，n阶马尔科夫链（Markov chain of order n）
# 齐夫定律：在一个自然语音语料库中，某个单词出现的频率与其在频率表中名次的常数次幂成反比.(较少数的单词会被经常使用到，绝大多数的单词很少被提及)

def read_time_machine():
    with open(r"D:\learn\深度学习\project\MNIST\article.txt", "r") as f:
        lines = f.readlines()
    return [re.sub("[^A-Za-z]+", " ", line).lower() for line in lines]


# 词元化
def tokenize(lines, token="word"):
    if token == "word":
        return [line.split() for line in lines]
    elif token == "char":
        return [list(line) for line in lines]
    else:
        print(f"error token:{token}")


def count_corpus(tokens):
    result = []
    if len(tokens) == 0 or isinstance(tokens[0], list):
        for token in tokens:
            for temp in token:
                result.append(temp)
    else:
        result = tokens
    return collections.Counter(result)


class Vocab:
    def __init__(self, tokens=None, min_freq=0, reserved_token=None):
        if tokens is None:
            tokens = []
        if reserved_token is None:
            reserved_token = []

        # 按照出现频率进行排序
        counter = count_corpus(tokens)
        self.token_freqs = sorted(counter.items(), key=lambda x: x[1], reverse=True)
        # 未知词元索引为0
        self.unk = 0
        self.uniq_tokens = ["<unk>"] + reserved_token
        self.uniq_tokens += [token for token, freq in self.token_freqs if
                             freq > min_freq and token not in reserved_token]
        self.idx_to_token, self.token_to_idx = [], dict()
        for token in self.uniq_tokens:
            self.idx_to_token.append(token)
            self.token_to_idx[token] = len(self.idx_to_token) - 1

    def __len__(self):
        return len(self.uniq_tokens)

    def __getitem__(self, tokens):
        if not isinstance(tokens, (list, tuple)):
            return self.token_to_idx.get(tokens, self.unk)
        return [self.__getitem__(token) for token in tokens]

    def to_tokens(self, indices):
        if not isinstance(indices, (list, tuple)):
            return self.idx_to_token[indices]
        return [self.to_tokens(index) for index in indices]


# 整合所有的功能
def load_corpus_time_machine(max_tokens=-1):
    lines = read_time_machine()
    tokens = tokenize(lines, "char")
    vocab = Vocab(tokens)
    corpus = [vocab[token] for line in tokens for token in line]
    if max_tokens > 0:
        corpus = corpus[:max_tokens]
    return corpus, vocab


tokens = tokenize(read_time_machine())
# print(tokens)
# 展平
corpus = [token for line in tokens for token in line]
# 词汇表
vacab = Vocab(corpus)
# print(vacab.token_freqs[:10])
# 取出所有的词频率
freqs = [freq for token, freq in vacab.token_freqs]
# 齐夫定律
d2l.plot(freqs, xlabel="token:x", ylabel="frequency:n(x)", xscale="log", yscale="log")
# d2l.plt.show()
# 二元语法 bi-grams
bigram_tokens = [pair for pair in zip(corpus[:-1], corpus[1:])]
# print(bigram_tokens[:10])
bigram_vocab = Vocab(bigram_tokens)
# print(bigram_vocab.token_freqs[:10])
bigram_freqs = [freq for token, freq in bigram_vocab.token_freqs]

trigram_tokens = [triple for triple in zip(corpus[:-2], corpus[1:-1], corpus[2:])]
trigram_vocab = Vocab(trigram_tokens)
# print(trigram_vocab.token_freqs[:10])
trigram_freqs = [freq for token, freq in trigram_vocab.token_freqs]
d2l.plot([freqs, bigram_freqs, trigram_freqs], xlabel="token:x", ylabel="frequency:n(x)", xscale="log", yscale="log",
         legend=["unigram", "bigram", "trigram"])


# d2l.plt.show()

# 序列数据取样
# 随机采样
def seq_data_iter_random(corpus, batch_size, num_steps):
    # 随机偏移量
    corpus = corpus[np.random.randint(0, num_steps - 1):-1]
    # 长度-1给标签数据留位置
    num_subseqs = (len(corpus) - 1) // num_steps
    # 子序列的起始索引
    initial_indices = list(range(0, num_subseqs * num_steps, num_steps))
    # np.random.shuffle会修改原始数据
    np.random.shuffle(initial_indices)

    def data(pos):
        return corpus[pos: pos + num_steps]

    num_batches = num_subseqs // batch_size
    for i in range(0, batch_size * num_batches, batch_size):
        initial_indices_per_batch = initial_indices[i: i + batch_size]
        # 取数据
        X = [data(j) for j in initial_indices_per_batch]
        Y = [data(j + 1) for j in initial_indices_per_batch]
        # yield是一个生成器generator，可用于迭代
        yield torch.tensor(X), torch.tensor(Y)


# corpus = np.arange(35)
# num_steps = 5
# num_subseqs = (len(corpus) - 1) // num_steps
# print(f"num_subseqs:{num_subseqs}")
# initial_indices = list(range(0, num_subseqs * num_steps,num_steps))
# print(initial_indices)
# np.random.shuffle会修改原始数据
# np.random.shuffle(initial_indices)
# print(initial_indices)
# my_seq = list(range(35))
# for X, Y in seq_data_iter_random(my_seq, batch_size=2, num_steps=5):
#    print(f"X:{X} \nY:{Y}")


# 顺序采样
def seq_data_iter_sequential(corpus, batch_size, num_steps):
    # 随机偏移量
    offset = np.random.randint(0, num_steps)
    # 有效token的长度
    num_tokens = ((len(corpus) - offset - 1) // num_steps) * num_steps
    Xs = torch.tensor(corpus[offset: num_tokens + offset])
    Ys = torch.tensor(corpus[offset + 1:num_tokens + 1 + offset])
    Xs, Ys = Xs.reshape(-1, num_steps), Ys.reshape(-1, num_steps)
    #    print(f"Xs:{Xs} \nYs:{Xs}")
    #    print(f"Xs:{Xs.shape} \nYs:{Xs.shape}")
    num_batches = Xs.shape[0] // batch_size
    for i in range(0, num_batches * batch_size, batch_size):
        X = Xs[i: i + batch_size, :]
        Y = Ys[i: i + batch_size, :]
        yield X, Y


my_seq = list(range(35))


# for X, Y in seq_data_iter_sequential(my_seq, 2, 5):
#    print(f"X:{X} \nY:{Y}")

# 把两个采样函数包装到类中，方便后续使用
class SeqDataLoader:
    def __init__(self, batch_size, num_steps, use_random_iter, max_tokens):
        if use_random_iter:
            self.data_iter_fn = seq_data_iter_random
        else:
            self.data_iter_fn = seq_data_iter_sequential

        self.corpus, self.vocab = load_corpus_time_machine(max_tokens)
        self.batch_size, self.num_steps = batch_size, num_steps

    def __iter__(self):
        return self.data_iter_fn(self.corpus, self.batch_size, self.num_steps)


# 返回数据迭代器和词汇表
def load_data_time_machine(batch_size, num_steps, user_random_iter=False, max_tokens=10000):
    data_iter = SeqDataLoader(batch_size, num_steps, user_random_iter, max_tokens)
    return data_iter, data_iter.vocab

# batch_size, num_steps = 2, 10
# loader, vocab = load_data_time_machine(batch_size, num_steps=num_steps)
# for i in loader:
#    print(i)
#    break
